#!/usr/bin/python

# a script to split a dataset list into 2 lists: 
# one list contains rootfile that has a trigger branch EF_e22vh_medium1
# and the other one dose not have it
# Lailin, Apr. 2012

import sys
ROOTSYS = '/afs/atlas.umich.edu/opt/root/lib'
sys.path.append(ROOTSYS)

from ROOT import TFile,TTree

import os
import glob

#input="test.list"
#outdir="new_list"
#treename="physics"

def main(input="test.list", treename="physics", outdir="new_list", **kw):

	try: fin=open(input, "r")
	except IOError:
		print "Error => $input can not be opened" % input
		sys.exit(-1)

	for line in fin.readlines():
		if line[-1]=="\n": line=line[:-1]
		if len(line)==0: continue
		dataset=line
		# get basename 
		sline=line.split("/")
		basename=''
		for i in range(len(sline)):
			if "mc11" in sline[i]: basename=sline[i]
			elif "data11" in sline[i]: basename=sline[i]
		if len(basename)==0:
			print "Warning => Neither MC nor Data, ingore this dataset: %s" % line
			continue
		# read root files	
		try: fdata=open(dataset, "r")
		except IOError:
			print "Warning => This dataset list can not be opened: %s" % dataset
			continue

		# split list: list1: has trigger branch, list2: not has trigger branch
		list_trig=[]
		list_notrig=[]
	
		for rootfile in fdata.readlines():
			if rootfile[-1]=="\n": rootfile=rootfile[:-1]
			print "INFO => Look into %s" % rootfile
			file=TFile(rootfile)
			if not file: 
				print "Warning => This root file can not be read in: %s" % rootfile
				continue
			tree=file.Get(treename)
			if not tree.FindBranch("EF_e22vh_medium1"):
				list_notrig.append(rootfile)
			else:
				list_trig.append(rootfile)
	
		fdata.close()
	
		# output
		if not os.path.isdir(outdir): os.mkdir(outdir)
		os.chdir(outdir)
	
		list0=basename.replace("list", "list0")
		list1=basename.replace("list", "list1")
		try: fout0=open(list0, "w")	
		except IOError:
			print "Warning => File %s Can not be written in to %s" % (list0, outdir)
			sys.exit(-1)
		try: fout1=open(list1, "w")	
		except IOError:
			print "Warning => File %s Can not be written in to %s" % (list1, outdir)
			sys.exit(-1)
	
		for file in list_notrig:
			fout0.write(file)
			fout0.write("\n")
		for file in list_trig:
			fout1.write(file)
			fout1.write("\n")
	
		fout0.close()
		fout1.close()
	
		os.chdir("../")
	
if __name__ == "__main__":
	argvv=sys.argv
	if len(argvv)==1: main(input="test.list", treename="physics", outdir="new_list")
	if len(argvv)==2: main(input=argvv[1])
	if len(argvv)==3: main(input=argvv[1], treename=argvv[2])
	if len(argvv)==4: main(input=argvv[1], treename=argvv[2], outdir=argvv[3])

